package com.su02.multi.chainrule;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

public abstract class Function implements NaiveFunction {
    private BitSet bs = new BitSet();
    private final Map<Integer, Function> backwardCache = new ConcurrentHashMap<>();

    @Override
    public Function backward(int id) {
        if (isNotRelated(id)) {
            return new Constant(0.);
        }
        return backwardCache.computeIfAbsent(id, this::doBackward);
    }

    protected abstract Function doBackward(int id);

    @Override
    public int size() {
        return bs.cardinality();
    }

    @Override
    public Function add(Function other) {
        if (isZero(this)) {
            return other;
        }
        if (isZero(other)) {
            return this;
        }
        return new Add(this, other);
    }

    @Override
    public Function sub(Function other) {
        return add(other.neg());
    }

    @Override
    public Function mul(Function other) {
        if (isZero(this) || isZero(other)) {
            return new Constant(0);
        }
        return new Multiply(this, other);
    }

    @Override
    public Function mul(double x) {
        if (x == 0.0) {
            return new Constant(0);
        }
        return mul(new Constant(x));
    }

    @Override
    public Function pow(double pow) {
        if (pow == 0.) {
            return new Constant(1);
        }
        if (pow == 1.) {
            return this;
        }
        return new Power(this, pow);
    }

    @Override
    public Function neg() {
        return new Neg(this);
    }

    protected boolean isNotRelated(int id) {
        return !bs.get(id);
    }

    protected void assertVariablesMatch(Map<Integer, Double> vector) {
        BitSet assertBs = getBs();
        Set<Integer> keys = new HashSet<>(vector.keySet());
        assertBs.stream().forEach(id -> {
            if (!keys.contains(id)) {
                throw new IllegalArgumentException();
            }
        });
    }

    public final BitSet getBs() {
        return bs;
    }

    protected void setBs(BitSet bs) {
        Objects.requireNonNull(bs);
        this.bs = bs;
    }

    private static boolean isZero(Function f) {
        return f instanceof Constant && ((Constant) f).getConstant() == 0;
    }
}